#include "calc_neighborhoods.h"

#include "../pdb_index/aa.h"
#include "aa_util.h"

#include <values.h>
#include <string.h>
#include <assert.h>

struct calc_neighborhoods * calc_neighborhoods_alloc(struct store * atomStore, 
						     struct mTree * atomTree,
						     double neighborhoodRadius,
						     char * pdbCode) {
  struct calc_neighborhoods * n = malloc(sizeof (struct calc_neighborhoods));
  
  n->filterState.count     = 3;

  n->filterState.functions = malloc(sizeof(void *(*)(void *, void *)) * 3);
  n->filterState.states    = malloc(sizeof(void*) * 3);

  n->filterState.functions[0] = (void*(*)(void*,void*))&filter_atom_by_name;
  n->filterState.states[0]    = " CA ";

  n->filterState.functions[1] = (void*(*)(void*,void*))&filter_atom_by_residue_number;
  n->filterState.states[1]    = 0;  // Set differently for each query...

  n->filterState.functions[2] = (void*(*)(void*,void*))&knn;
  n->filterState.states[2]    = knn_state_alloc(NEIGHBORS, 0, (double(*)(const void*,const void*))&atom_distance, sizeof(struct atom));  

  n->atomStore = atomStore;
  n->atomTree  = atomTree;

  n->neighborhoodRadius = neighborhoodRadius;

  strncpy(n->pdbCode, pdbCode, 4);

  n->storePosition = 1; // The first position (0) in the store is a placeholder.



  return n;
}
						    
void calc_neighborhoods_free(struct calc_neighborhoods * n) {
  knn_state_free(n->filterState.states[2]);
  free (n->filterState.functions);
  free (n->filterState.states);
  free (n);
}

int calc_neighborhoods_next(struct calc_neighborhoods * n, struct aa_bond_neighborhood * neighborhood) {
  for( ; n->storePosition != -1; n->storePosition++) {
    struct atom * atom = (struct atom*)store_read(n->atomStore, n->storePosition);

    assert(n->storePosition != 0);

    if(atom->serialNumber == 0) {
      store_release(n->atomStore, n->storePosition);

      n->storePosition = -1;

      return 0;
    } else if(!strncmp(atom->name, " CA ", 4)) {
      off_t j = find_in_aa_by_name(n->atomStore, n->storePosition, " CB ");
      off_t k = find_in_aa_by_name(n->atomStore, n->storePosition, " C  ");
      struct aa aa;

      if(j >= 0 && k >= 0) {
	struct atom * atm;
	struct aa neighbors[NEIGHBORS];
	
	int c;

	// Build center amino acid

	atm = (struct atom*)store_read(n->atomStore, j); // j is the CB index
	aa_init(&aa, atom, atm);
	store_release(n->atomStore, j);
	
	/* Find the rest of the neighborhood */
	
	// First, set up the query filters.
	// Residue filter needs to know which residue to ignore:
        n->filterState.states[1] = (void*)atom->residue;
	// K-Nearest Neighbors needs to know the query object:
	((struct knn_state*)n->filterState.states[2])->queryObject = atom;

	// Run the query

	mTree_rangeQuery(n->atomTree, atom, n->neighborhoodRadius, (void(*)(void*,void*))&filterAnd, &(n->filterState));

	// The knn state now contains all of the nearest neighbors--copy them into the local neighbors array.

	for(c = 0; c < NEIGHBORS; c++) {
	  struct knn_state* knn_state = (struct knn_state*)n->filterState.states[2];
	  if(knn_state->results_distance[c] != MAXDOUBLE) {
	    struct atom * neighbor_atom = &((struct atom*)knn_state->results)[c];
	    struct atom * neighbor_cb;
	    off_t l = find_in_aa_by_name(n->atomStore, neighbor_atom->serialNumber, " CB ");

	    if(l >= 0) {
	      neighbor_cb = store_read(n->atomStore, l);
	      aa_init(&neighbors[c], neighbor_atom, neighbor_cb);
	      store_release(n->atomStore, l);
	    } else {
	      aa_init(&neighbors[c], neighbor_atom, neighbor_atom);
	    }
	  }
	}

	// Done with knn_state:
	knn_state_reset((struct knn_state*)n->filterState.states[2]);
	
	/* k is the backbone carbon, which we need to pass into the neighborhood initialization function. */

	atm = (struct atom*) store_read(n->atomStore, k);
	aa_bond_neighborhood_init(neighborhood, n->pdbCode, &aa, atm, neighbors);
	
	store_release(n->atomStore, k);
	store_release(n->atomStore, n->storePosition);
	n->storePosition++;
	return 1;

      }
    } 
    store_release(n->atomStore, n->storePosition);
  }
  
  return 0;
}

void calc_neighborhoods(char * pdbCode, double neighborhoodRadius,
			struct store * atomStore, struct mTree * atomTree,
			struct store * neighStore, struct mTree * neighTree) {
  struct calc_neighborhoods * cn = calc_neighborhoods_alloc(atomStore, atomTree, neighborhoodRadius, pdbCode);
  struct aa_bond_neighborhood tmp;
  while(calc_neighborhoods_next(cn, &tmp)) {
    off_t i = store_new(neighStore);
    struct aa_bond_neighborhood * n = store_read(neighStore, i);
    
    //    aa_bond_neighborhood_print(&tmp);

    memcpy(n, &tmp, sizeof(struct aa_bond_neighborhood));

    store_write(neighStore, i);
    mTree_insert(neighTree, i);
    store_release(neighStore, i);
  }
  
  calc_neighborhoods_free(cn);
}
/*
int main(int argc, char** argv) {
  struct store * atomStore, * neighStore;
  struct mTree * atomTree,  * neighTree;
  
  if(argc != 6) {
    printf("This program takes exactly five arguments; the four letter pdb code, the neighborhood store, the neighborhood store fat, the neighborhood index, and its fat.  (Argument count was %d)  Size of aa_bond_neighborhood is %d\n", argc - 1, sizeof(struct aa_bond_neighborhood));
    return 1;
  }
  
  atomStore  = store_open_cached("neigh_store.atoms", "neigh_store.atoms.fat", sizeof(struct atom), 100, 1000000);
  atomTree   = mTree_open_cached(atomStore, "neigh_store.mTree", "neigh_store.mTree.fat", 120000, (double(*)(void*,void*))&atom_distance);
  neighStore = store_open_cached(argv[2],argv[3], sizeof(struct aa_bond_neighborhood), 100, 1000);
  neighTree  = mTree_open_cached(neighStore, argv[4], argv[5], 10000, (double(*)(void*,void*))&aa_bond_neighborhood_distance);

  calc_neighborhoods(argv[1], 30 * 1000, atomStore, atomTree, neighStore, neighTree);

  mTree_close(neighTree);
  store_close(neighStore);

  mTree_close(atomTree);
  store_close(atomStore);

  return 0;
  
}
*/

/*
int main_old(int argc, char** argv) {

  struct store * atomStore = store_open_cached("neigh_store.atoms", "neigh_store.atoms.fat", 
					       sizeof(struct atom), 100, 100000);
  struct mTree * atomTree  = mTree_open_cached(atomStore, "neigh_store.mTree", "neigh_store.mTree.fat", 
					       120000, (double(*)(void*,void*))&atom_distance);
  void* (*functions[])(void*,void*) = {(void*(*)(void*,void*))&filter_atom_by_name, 
				       (void*(*)(void*,void*))&filter_atom_by_residue_number,
				       (void*(*)(void*,void*))&knn};
  struct knn_state * knn_state = knn_state_alloc(NEIGHBORS, 0, (double(*)(void*,void*))&atom_distance, sizeof(struct atom));
  
  struct store * neighborhoodStore;

  struct mTree * neighborhoodTree;

  void* states[] = {" CA ",
		    0,//The residue of the query atom-set below... 
		    0 //knn_state-see below
		    
                   };
  struct filter_state filter_state = {3, functions, states};

  off_t i = 0;
  
  int done = 0;


  if(argc != 6) {
    printf("This program takes exactly five arguments; the four letter pdb code, the neighborhood store, the neighborhood store fat, the neighborhood index, and its fat.  (Argument count was %d)  Size of aa_bond_neighborhood is %d\n", argc - 1, sizeof(struct aa_bond_neighborhood));
    return 1;
  }

  neighborhoodStore = store_open_cached(argv[2], argv[3], sizeof(struct aa_bond_neighborhood), 100, 1000);
  neighborhoodTree  = mTree_open_cached(neighborhoodStore, argv[4], argv[5], 10000, (double(*)(void*,void*))&aa_bond_neighborhood_distance);

  states[2] = knn_state;
  for(i = 0; !done; i++) {
    struct atom * atom = (struct atom*)store_read(atomStore, i);
    if(atom->residue == 0) { 
      //PDB counts residues starting at 1, and the store zeros newly
      //initialized space, so we're at the end of the store.
      done = 1;
    } else if(!strncmp(atom->name, " CA ", 4)) {
      off_t j = find_in_aa_by_name(atomStore, i, " CB ");
      off_t k = find_in_aa_by_name(atomStore, i, " C  ");
      struct aa aa;
      //      printf("Checking %d\n", (int)i);
      if(j >= 0 && k >= 0) {
	struct atom * atm;
	struct aa neighbors[NEIGHBORS];
	off_t neighborhood_idx = store_new(neighborhoodStore);
	struct aa_bond_neighborhood * neighborhood = store_read(neighborhoodStore, neighborhood_idx);
	int n;

	// Build the center amino acid.

	atm = (struct atom*)store_read(atomStore, j); // j is the CB index
	aa_init(&aa, atom, atm);
	store_release(atomStore, j);

	// Find the rest of the neighborhood 

	// Initialize our query filters. 

	// K-Nearest Neighbors needs to have the query object...
	knn_state->queryObject = atom;
	// Residue filter needs to know which residue to ignore...
	states[1] = (void*)atom->residue;  

	// Run a query with a 30 angstrom radius.
	mTree_rangeQuery(atomTree, atom, 30 * 1000, (void(*)(void*,void*))&filterAnd, &filter_state);  //10 angstroms * 1000 milliA / A

	// knn now contains all of the nearest neighbors--copy them into
	// our local neighbors array.
	
	for(n = 0; n<NEIGHBORS; n++) {
	  if(knn_state->results_distance[n] != MAXDOUBLE) {
	    struct atom * neighbor_atom = &((struct atom*)knn_state->results)[n];
	    off_t l = find_in_aa_by_name(atomStore, neighbor_atom->serialNumber, " CB ");
	    struct atom * neighbor_cb;
	    if(l >= 0) {
	      neighbor_cb = store_read(atomStore, l);
	      aa_init(&neighbors[n], neighbor_atom, neighbor_cb);
	      store_release(atomStore, l);
	    } else {
	      aa_init(&neighbors[n], neighbor_atom, neighbor_atom);
	    }
	  }
	}

	//	knn_state_print(knn_state, (void(*)(void*,void*))atom_print);


	// Don't need knn until the next loop iteration.
	knn_state_reset(knn_state);


	// k is the backbone carbon, which we need to pass into 
	//   the neighborhood initialization function. 


	atm = (struct atom*)store_read(atomStore, k);
	aa_bond_neighborhood_init(neighborhood, argv[1], &aa, atm, neighbors);
	store_release(atomStore, k);

	store_write(neighborhoodStore, neighborhood_idx);

	mTree_insert(neighborhoodTree, neighborhood_idx);

	store_release(neighborhoodStore, neighborhood_idx);
	

	//	aa_bond_neighborhood_print(&neighborhood);

      }
    }
    store_release(atomStore, i);
  }
  knn_state_free(knn_state);
  mTree_close(atomTree);
  store_close(atomStore);
  mTree_close(neighborhoodTree);
  store_close(neighborhoodStore);
  return 0;
}
*/
